# from ultralytics import YOLO
import numpy as np
from typing import Optional
import queue
import threading
import time
from concurrent.futures import ThreadPoolExecutor
from collections import deque
from image_provider import ImageProvider
from compute_position import ComputePosition
from opt_visualizer import MergedPosePositionVisualizer
from seg_pose_rk import SegPoseRK  # 这是我们之前实现的SegPose类
from performance_monitor import PerformanceMonitor
from utils import *

logger = setup_logger()


def main():
    # 初始化SegPose处理器 - 增加队列大小避免阻塞读取
    processor = SegPoseRK(
        seg_model_path="./models/merge_yolo11m_seg_int8_rk3588.rknn",
        pose_model_path="./models/merge_yolo11n_pose_no_rk3588.rknn",
        seg_workers=2,
        pose_workers=2,
        seg_npus=[1, 2],
        pose_npus=[0],
        target_fps=30
    )
    processor.start()
    
    # 初始化图像提供者
    image_provider = ImageProvider(from_folder=True)
    
    compute_position = ComputePosition()

    st_pose, st_vector = setup_standard_position(None, compute_position)

    visualizer = MergedPosePositionVisualizer(st_pose=st_pose, st_vector=st_vector)
    
    # 初始化性能监控（现在是Thread类，传入logger）
    perf_monitor = PerformanceMonitor(update_interval=2.0, logger=logger)
    
    # 创建线程间通信队列 - 增加队列大小
    result_processing_queue = queue.Queue(maxsize=50)
    visualization_queue = queue.Queue(maxsize=20)
    
    # 全局退出标志
    global_shutdown_flag = threading.Event()
    
    # 线程1：图像采集和提交
    def image_acquisition_thread():
        """图像采集和提交线程 - 严格按时间间隔读取"""
        target_interval = 1.0 / processor.target_fps  # 使用处理器的目标FPS
        frame_count = 0
        dropped_frames = 0
        
        try:
            # 记录开始时间，用于精确控制时间间隔
            start_time = time.time()
            
            while not global_shutdown_flag.is_set():
                # 计算当前帧应该开始的时间
                target_time = start_time + frame_count * target_interval
                current_time = time.time()
                
                # 如果还没到时间，就等待
                if current_time < target_time:
                    time.sleep(target_time - current_time)
                
                frame_start = time.time()
                
                # 获取图像
                img_data = image_provider.get_next_image()
                if img_data is None:
                    logger.info("Image provider finished.")
                    break
                
                # 统计读图片（无论是否成功提交都要统计，因为图片确实被读取了）
                perf_monitor.update_stats('read')
                frame_count += 1
                
                # 尝试提交到处理器（非阻塞）
                task_id = processor.put(img_data)
                if task_id == -1:
                    dropped_frames += 1
                    if dropped_frames % 10 == 1:  # 每10帧丢弃输出一次警告
                        logger.warning(f"Processing queue full, frame dropped (total dropped: {dropped_frames})")
                else:
                    logger.debug(f"Frame {frame_count} submitted (ID: {task_id})")
                
                # 记录图像读取和提交的总耗时
                read_submit_time = time.time() - frame_start
                if read_submit_time > target_interval * 0.8:  # 如果超过目标间隔的80%
                    logger.warning(f"Image read+submit took {read_submit_time*1000:.1f}ms (target: {target_interval*1000:.1f}ms)")
        
        except Exception as e:
            logger.error(f"Image acquisition thread error: {e}")
        finally:
            # 图像采集结束，触发关闭流程
            logger.info(f"Image acquisition completed. Total frames: {frame_count}, Dropped: {dropped_frames}")
            global_shutdown_flag.set()
    
    # 线程2：结果获取
    def result_fetch_thread():
        """结果获取线程 - 高频获取结果"""
        fetch_interval = 0.01  # 10ms 更频繁地获取结果
        results_fetched = 0
        
        try:
            while not global_shutdown_flag.is_set():
                start_time = time.time()
                
                # 批量获取结果，提高效率
                batch_results = []
                for _ in range(5):  # 每次最多获取5个结果
                    result = processor.get_result(block=False)
                    if result is not None:
                        batch_results.append(result)
                    else:
                        break
                
                # 处理获取到的结果
                for result in batch_results:
                    results_fetched += 1
                    
                    # 检查是否为有效检测结果
                    has_valid_detection = (
                        result.get('seg_result') is not None and 
                        result.get('pose_result') is not None and
                        result.get('seg_success', False) and
                        result.get('pose_success', False)
                    )
                    
                    # 统计检测成功（只有成功的才计入）
                    if has_valid_detection:
                        perf_monitor.update_stats('detection', success=True)
                    
                    # 将结果放入处理队列
                    try:
                        result_processing_queue.put(result, block=False)
                    except queue.Full:
                        logger.warning("Processing queue full, result dropped")
                
                if len(batch_results) > 0:
                    logger.debug(f"Fetched {len(batch_results)} results (total: {results_fetched})")
                
                # 等待剩余时间
                elapsed = time.time() - start_time
                if elapsed < fetch_interval:
                    time.sleep(fetch_interval - elapsed)
            
            # 处理剩余结果
            remaining_count = 0
            while processor.pending_count() > 0:
                result = processor.get_result(block=False)
                if result is not None:
                    remaining_count += 1
                    has_valid_detection = (
                        result.get('seg_result') is not None and 
                        result.get('pose_result') is not None and
                        result.get('seg_success', False) and
                        result.get('pose_success', False)
                    )
                    if has_valid_detection:
                        perf_monitor.update_stats('detection', success=True)
                    try:
                        result_processing_queue.put(result, block=False)
                    except queue.Full:
                        logger.warning("Processing queue full, result dropped")
                else:
                    time.sleep(0.01)
            
            if remaining_count > 0:
                logger.info(f"Processed {remaining_count} remaining results")
                    
        except Exception as e:
            logger.error(f"Result fetch thread error: {e}")
    
    # 线程3：结果处理（计算pose和vector）
    def result_processing_thread():
        """结果处理线程"""
        try:
            while not global_shutdown_flag.is_set() or not result_processing_queue.empty():
                try:
                    # 阻塞获取结果（带超时以检查退出条件）
                    result = result_processing_queue.get(timeout=0.5)
                    logger.debug(f"Processing result for task {result.get('task_id', 'unknown')}")
                    
                    pose = None
                    vector = None
                    # print("result:",result)
                    # 检查结果的有效性
                    if (result.get('image_data') is not None and 
                        result.get('seg_result') is not None and 
                        result.get('pose_result') is not None):
                        
                        time_stamp = result['image_data'].timestamp
                        color_img = result['image_data'].color_img
                        depth_img = result['image_data'].depth_img
                        seg_result = result['seg_result'][0][0]
                        pose_result = result['pose_result'][0] 
                        
                        print("seg_result:",seg_result)
                        print("pose_result:",pose_result)
                        # logger.debug(f"Processing valid result - seg_result type: {type(seg_result)}, pose_result type: {type(pose_result)}")
                        
                        try:
                            pose, vector = compute_position.process_image_pair(
                                time_stamp, color_img, depth_img, seg_result, pose_result
                            )
                            print("pose,vector:",pose,vector)
                            logger.debug(f"Position computation completed - pose: {pose is not None}, vector: {vector is not None}")
                        except Exception as compute_error:
                            logger.error(f"Position computation error: {compute_error}")
                            pose, vector = None, None
                    else:
                        logger.debug("Invalid or incomplete result, skipping position computation")

                    # 检查位置计算是否成功
                    position_success = pose is not None and vector is not None
                    if position_success:
                        perf_monitor.update_stats('position', success=True)
                        logger.debug("Position computation successful")
                    else:
                        logger.debug("Position computation failed or skipped")
                    
                    # # 始终放入可视化队列，即使pose或vector为None
                    # visualization_data = {
                    #     'pose': pose,  # 可能为None
                    #     'vector': vector,  # 可能为None
                    #     'color_img': result['image_data'].color_img if result.get('image_data') else None,
                    #     'timestamp': result['image_data'].timestamp if result.get('image_data') else time.time(),
                    #     'has_detection': position_success
                    # }
                    
                    # try:
                    #     visualization_queue.put(visualization_data, block=False)
                    #     if position_success:
                    #         logger.debug("Data with detection sent to visualization queue")
                    #     else:
                    #         logger.debug("Data without detection sent to visualization queue")
                    # except queue.Full:
                    #     # 可视化队列满时，丢弃最旧的数据
                    #     try:
                    #         visualization_queue.get(block=False)
                    #         visualization_queue.put(visualization_data, block=False)
                    #         logger.warning("Visualization queue full, dropped old data")
                    #     except queue.Full:
                    #         logger.warning("Visualization queue full, current data dropped")
                    
                    # 标记任务完成
                    result_processing_queue.task_done()
                    
                except queue.Empty:
                    continue
                except Exception as process_error:
                    logger.error(f"Result processing error: {process_error}")
                    try:
                        result_processing_queue.task_done()
                    except:
                        pass
                    
        except Exception as e:
            logger.error(f"Result processing thread error: {e}")
    
    # 线程4：可视化线程
    def visualization_thread():
        """可视化处理线程"""
        try:
            vis_timeout_count = 0
            max_timeout_count = 10
            
            while not global_shutdown_flag.is_set() or not visualization_queue.empty():
                try:
                    # 阻塞获取可视化数据（带超时以检查退出条件）
                    vis_data = visualization_queue.get(timeout=0.5)
                    vis_timeout_count = 0  # 重置超时计数
                    
                    # 执行可视化
                    try:
                        start_vis = time.time()
                        
                        # 获取当前可视化FPS
                        current_fps = perf_monitor.get_latest_fps('visualization')
                        
                        # 可视化 - 添加超时保护
                        try:
                            result_path, canvas = visualizer.merge_visualize(
                                vis_data['pose'],           # pose_dict
                                vis_data['vector'],         # test_vector  
                                vis_data['color_img'],      # color_img
                                timestamp=vis_data['timestamp'],
                                display=True,
                                save_async=False,
                                fps=current_fps,            # 传入FPS
                                frame_count=perf_monitor.visualization_count  # 传入计数
                            )
                            
                            # 统计可视化完成
                            perf_monitor.update_stats('visualization')
                            
                            vis_time = time.time() - start_vis
                            logger.debug(f"Visualization completed in {vis_time*1000:.1f}ms")
                            
                        except Exception as vis_error:
                            logger.error(f"Visualization merge_visualize error: {vis_error}")
                            # 即使可视化失败也要统计，避免阻塞
                            perf_monitor.update_stats('visualization')
                        
                    except Exception as e:
                        logger.error(f"Visualization wrapper error: {e}")
                        # 确保统计更新
                        perf_monitor.update_stats('visualization')
                    
                    # 标记可视化任务完成
                    visualization_queue.task_done()
                    
                except queue.Empty:
                    vis_timeout_count += 1
                    # 如果长时间没有数据，检查是否应该退出
                    if vis_timeout_count >= max_timeout_count:
                        logger.info("Visualization thread: No data for extended period, checking exit condition")
                        if global_shutdown_flag.is_set() and result_processing_queue.empty():
                            logger.info("Visualization thread: Exit conditions met")
                            break
                        vis_timeout_count = 0  # 重置计数器继续等待
                    continue
                    
        except Exception as e:
            logger.error(f"Visualization thread error: {e}")
        finally:
            logger.info("Visualization thread cleanup complete")
    
    # 创建并启动所有线程
    acquisition_thread = threading.Thread(target=image_acquisition_thread, name="ImageAcquisition")
    fetch_thread = threading.Thread(target=result_fetch_thread, name="ResultFetch")
    processing_thread = threading.Thread(target=result_processing_thread, name="ResultProcessing")
    # vis_thread = threading.Thread(target=visualization_thread, name="Visualization")
    
    # 性能监控现在是Thread对象，直接启动
    threads = [acquisition_thread, fetch_thread, processing_thread, perf_monitor]
    
    # 启动所有线程
    for thread in threads:
        thread.start()
        logger.info(f"Started thread: {thread.name}")
    
    try:
        # 等待图像采集线程结束（这会触发shutdown流程）
        acquisition_thread.join()
        logger.info("Image acquisition finished. Initiating shutdown sequence...")
        
        # 等待其他处理线程结束
        fetch_thread.join()
        logger.info("Result fetching finished.")
        
        processing_thread.join()
        logger.info("Result processing finished.")
        
        # vis_thread.join()
        # logger.info("Visualization finished.")
        
        # 关闭性能监控线程
        perf_monitor.shutdown()
        perf_monitor.join()
        logger.info("Performance monitoring finished.")
        
        # 输出最终统计信息
        final_status = processor.get_status()
        logger.info(f"Final processor status: {final_status}")
        
    except KeyboardInterrupt:
        logger.info("Received interrupt signal. Shutting down...")
        global_shutdown_flag.set()
        
        # 关闭性能监控
        perf_monitor.shutdown()
        
        # 等待所有线程结束
        for thread in threads:
            if thread.is_alive():
                thread.join(timeout=5.0)  # 最多等待5秒
                if thread.is_alive():
                    logger.warning(f"Thread {thread.name} did not finish within timeout")
    
    except Exception as e:
        logger.error(f"Main thread error: {e}")
        global_shutdown_flag.set()
    
    finally:
        # 确保性能监控关闭
        if hasattr(perf_monitor, 'shutdown_flag') and not perf_monitor.shutdown_flag:
            perf_monitor.shutdown()
        
        # 正式关闭处理器
        processor.stop()  # 使用正确的方法名
        logger.info("System shutdown complete.")

if __name__ == "__main__":
    main()